package com.sz.admin.ai.util;

import com.sz.admin.ai.factory.embeddingFactory.entity.EmbeddingModelEntity;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.ai.vectorstore.pgvector.PgVectorStore;
import org.springframework.jdbc.core.JdbcTemplate;
import org.springframework.stereotype.Component;

/**
 * @描述:
 * @Author: TangYuan
 * @Date: 2025/2/25 11:02
 */
@Component
public class PgVectorStoreUtil {
    
    public VectorStore getPgVectorStore(JdbcTemplate jdbcTemplate, EmbeddingModel embeddingModel, EmbeddingModelEntity modelEntity) {
        // PGVectorStore创建
        PgVectorStore vectorStore = PgVectorStore
                .builder(jdbcTemplate, embeddingModel)
                .schemaName("public")
                .vectorTableName("vector_store_" + modelEntity.getModelDimension())
                .vectorTableValidationsEnabled(PgVectorStore.DEFAULT_SCHEMA_VALIDATION)
                .dimensions(modelEntity.getModelDimension())
                .distanceType(PgVectorStore.PgDistanceType.COSINE_DISTANCE)
                .removeExistingVectorStoreTable(false)
                .indexType(PgVectorStore.PgIndexType.HNSW)
                .initializeSchema(true)
                .maxDocumentBatchSize(PgVectorStore.MAX_DOCUMENT_BATCH_SIZE)
                .build();
        
        vectorStore.afterPropertiesSet();
        return vectorStore;
    }
}
